import asyncio
from temporalio import workflow
from temporalio.client import Client
import psycopg2

# 工作流接口
class GaussDBWorkflow:
    async def create_record(self, record_data):
        """创建数据库记录的抽象方法"""
        raise NotImplementedError

    async def read_record(self, record_id):
        """读取数据库记录的抽象方法"""
        raise NotImplementedError

    async def update_record(self, record_id, updated_data):
        """更新数据库记录的抽象方法"""
        raise NotImplementedError

    async def delete_record(self, record_id):
        """删除数据库记录的抽象方法"""
        raise NotImplementedError

# 工作流具体实现类
@workflow.defn
class GaussDBWorkflowImpl(GaussDBWorkflow):
    def __init__(self):
        self.conn = None

    # 工作流启动时建立数据库连接
    @workflow.run
    async def run(self):
        try:
            self.conn = psycopg2.connect(
                host="xxx.xxx.xxx.xxx",
                port="xxxx",
                database="xxxx",
                user="xxxx",
                password="xxxxxx"
            )
            await self.setup_table()
        except psycopg2.Error as e:
            print(f"数据库连接或初始化表失败: {e}")

    async def setup_table(self):
        """创建示例数据表（如果不存在）"""
        with self.conn.cursor() as cursor:
            try:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS xjw.example_table (
                        id SERIAL PRIMARY KEY,
                        name VARCHAR(100),
                        age INT
                    );
                """)
                self.conn.commit()
            except psycopg2.Error as e:
                print(f"创建表失败: {e}")

    async def create_record(self, record_data):
        """向表中创建记录"""
        with self.conn.cursor() as cursor:
            try:
                cursor.execute(
                    "INSERT INTO xjw.example_table (name, age) VALUES (%s, %s)",
                    (record_data['name'], record_data['age'])
                )
                self.conn.commit()
            except psycopg2.Error as e:
                print(f"创建记录失败: {e}")

    async def read_record(self, record_id):
        """从表中读取记录"""
        with self.conn.cursor() as cursor:
            try:
                cursor.execute("SELECT * FROM xjw.example_table WHERE id = %s", (record_id,))
                result = cursor.fetchone()
                return result
            except psycopg2.Error as e:
                print(f"读取记录失败: {e}")
                return None

    async def update_record(self, record_id, updated_data):
        """更新表中的记录"""
        with self.conn.cursor() as cursor:
            try:
                cursor.execute(
                    "UPDATE xjw.example_table SET name = %s, age = %s WHERE id = %s",
                    (updated_data['name'], updated_data['age'], record_id)
                )
                self.conn.commit()
            except psycopg2.Error as e:
                print(f"更新记录失败: {e}")

    async def delete_record(self, record_id):
        """从表中删除记录"""
        with self.conn.cursor() as cursor:
            try:
                cursor.execute("DELETE FROM xjw.example_table WHERE id = %s", (record_id,))
                self.conn.commit()
            except psycopg2.Error as e:
                print(f"删除记录失败: {e}")

    # 工作流结束时关闭数据库连接
    @workflow.signal
    async def close_connection(self):
        if self.conn:
            self.conn.close()

async def main():
    client = Client("xxx.xxx.xxx.xxx:xxxx")
    # 启动工作流，指定工作流类型、ID以及任务队列
    workflow_options = await client.start_workflow(
        GaussDBWorkflowImpl,
        id="gauss-db-workflow-1",
        task_queue="gauss-db-task-queue"
    )
    workflow_instance = client.get_workflow_handle(
        workflow_id="gauss-db-workflow-1",
        task_queue="gauss-db-task-queue"
    )

    # 创建记录
    create_data = {'name': 'John Doe', 'age': 30}
    create_data1 = {'name': 'Ho Hui Mei', 'age': 20}
    await workflow_instance.create_record(create_data)
    await workflow_instance.create_record(create_data1)

    # 读取记录
    record_id = 1  # 假设刚插入的记录id为1，实际中根据插入情况获取正确id
    record = await workflow_instance.read_record(record_id)
    print(f"读取的记录: {record}")

    # 更新记录
    update_data = {'name': 'Updated John Doe', 'age': 31}
    await workflow_instance.update_record(record_id, update_data)

    # 再次读取记录查看更新情况
    updated_record = await workflow_instance.read_record(record_id)
    print(f"更新后的记录: {updated_record}")

    # 删除记录
    await workflow_instance.delete_record(record_id)
    final_record = await workflow_instance.read_record(record_id)
    print(f"删除后再次读取记录（应为None）: {final_record}")

    # 关闭工作流连接
    await workflow_instance.close_connection()

if __name__ == "__main__":
    asyncio.run(main())
